import os

from langchain.agents import create_structured_chat_agent, AgentExecutor
from langchain.memory import ConversationBufferMemory
from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.utilities import SQLDatabase
from langchain_core.tools import tool, Tool
from langchain_openai import ChatOpenAI
from langchain import hub

os.environ["OPENWEATHERMAP_API_KEY"] = "879544e7845335ee2f4d3168fd81cd55"

db = SQLDatabase.from_uri("sqlite:///../my_db.db")

key = 'sk-7wnDma9l5GuVbq38B3C07f50290147148a0809B117A1C1Ad'
model = ChatOpenAI(model="gpt-3.5-turbo",
                   openai_api_key=key,
                   openai_api_base="https://api.aigc369.com/v1",
                   temperature=0)

db_agent_executor = create_sql_agent(
    db=db, llm=model, agent_type="openai-tools", verbose=True, handle_parsing_errors=True
)


@tool
def get_user_info(user_id: int) -> str:
    """获取用户信息的工具，当你需要获取用户信息时，请使用此工具"""
    user_info = db_agent_executor.invoke({"input": f"SELECT * FROM user WHERE user_id = {user_id}"})
    department_info = db_agent_executor.invoke({"input": f"SELECT * FROM department WHERE department_id = (SELECT department_id FROM user WHERE user_id = {user_id})"})
    performance_info = db_agent_executor.invoke({"input": f"SELECT * FROM performance WHERE performance_id = (SELECT performance_id FROM user WHERE user_id = {user_id})"})
    return user_info, department_info, performance_info


def create_react_agent():
    tools = [
        get_user_info,
        Tool(
            name="数据库查询工具",
            description="""
                当你本身无法获取到问题追定的数据信息的时候，你可以使用此数据库查询工具，将用户的问题转为SQL后，进行查询操作，最后工具会返回一个数据查询结果
            """,
            func=db_agent_executor.invoke
        )
    ]

    prompt = hub.pull("hwchase17/structured-chat-agent")

    agent = create_structured_chat_agent(
        llm=model,
        prompt=prompt,
        tools=tools
    )

    memory = ConversationBufferMemory(
        return_messages=True,
        memory_key="chat_history"
    )

    agent_executor = AgentExecutor.from_agent_and_tools(
        agent=agent,
        tools=tools,
        memory=memory,
        verbose=True,
        handle_parsing_errors=True
    )

    return agent_executor


react_agent = create_react_agent()

while True:
    input_text = input("请输入你的问题:")
    print(react_agent.invoke({"input": input_text}))
